lib/Mi manual de algoritmos/version_actual/src/estructuras/segment_tree.cpp

   1 class SegmentTree{
   2 public:
   3   vector<int> arr, tree;
   4   int n;
   5
   6   SegmentTree(){}
   7   SegmentTree(const vector<int> &arr) : arr(arr) {
   8     initialize();
   9   }
  10
  11   //must be called after assigning a new arr.
  12   void initialize(){
  13     n = arr.size();
  14     tree.resize(4*n + 1);
  15     initialize(0, 0, n-1);
  16   }
  17
  18   int query(int query_left, int query_right) const{
  19     return query(0, 0, n-1, query_left, query_right);
  20   }
  21
  22   void update(int where, int what){
  23     update(0, 0, n-1, where, what);
  24   }
  25
  26 private:
  27   int initialize(int node, int node_left, int node_right);
  28   int query(int node, int node_left, int node_right,
  29             int query_left, int query_right) const;
  30   void update(int node, int node_left, int node_right,
  31               int where, int what);
  32 };
  33
  34 int SegmentTree::initialize(int node,
  35                             int node_left, int node_right){
  36   if (node_left == node_right){
  37     tree[node] = node_left;
  38     return tree[node];
  39   }
  40   int half = (node_left + node_right) / 2;
  41   int ans_left = initialize(2*node+1, node_left, half);
  42   int ans_right = initialize(2*node+2, half+1, node_right);
  43
  44   if (arr[ans_left] <= arr[ans_right]){
  45     tree[node] = ans_left;
  46   }else{
  47     tree[node] = ans_right;
  48   }
  49   return tree[node];
  50 }
  51
  52 int SegmentTree::query(int node, int node_left, int node_right,
  53                        int query_left, int query_right) const{
  54   if (node_right < query_left || query_right < node_left)
  55     return -1;
  56   if (query_left <= node_left && node_right <= query_right)
  57     return tree[node];
  58
  59   int half = (node_left + node_right) / 2;
  60   int ans_left = query(2*node+1, node_left, half,
  61                        query_left, query_right);
  62   int ans_right = query(2*node+2, half+1, node_right,
  63                         query_left, query_right);
  64
  65   if (ans_left == -1) return ans_right;
  66   if (ans_right == -1) return ans_left;
  67
  68   return (arr[ans_left] <= arr[ans_right] ? ans_left : ans_right);
  69 }
  70
  71 void SegmentTree::update(int node, int node_left, int node_right,
  72                          int where, int what){
  73   if (where < node_left || node_right < where) return;
  74   if (node_left == where && where == node_right){
  75     arr[where] = what;
  76     tree[node] = where;
  77     return;
  78   }
  79   int half = (node_left + node_right) / 2;
  80   if (where <= half){
  81     update(2*node+1, node_left, half, where, what);
  82   }else{
  83     update(2*node+2, half+1, node_right, where, what);
  84   }
  85   if (arr[tree[2*node+1]] <= arr[tree[2*node+2]]){
  86     tree[node] = tree[2*node+1];
  87   }else{
  88     tree[node] = tree[2*node+2];
  89   }
  90 }